import os
import torch
from torch.utils.data import Dataset
from PIL import Image, ImageDraw
import pandas as pd
from torchvision import datasets, transforms

class CLIPDataset(Dataset):
    """CLIP数据集类"""
    
    def __init__(self, image_dir, text_file, processor, transform=None, max_length=77):
        """
        初始化数据集
        
        Args:
            image_dir: 图像目录路径
            text_file: 文本描述文件路径
            processor: CLIP处理器
            transform: 图像变换
            max_length: 文本最大长度
        """
        self.image_dir = image_dir
        self.processor = processor
        self.transform = transform
        self.max_length = max_length
        
        # 加载文本描述
        if os.path.exists(text_file):
            if text_file.endswith('.csv'):
                self.data = pd.read_csv(text_file)
            else:
                # 假设是每行包含图像文件名和描述的文本文件
                self.data = []
                with open(text_file, 'r', encoding='utf-8') as f:
                    for line in f:
                        parts = line.strip().split('\t')
                        if len(parts) >= 2:
                            self.data.append({
                                'image_name': parts[0],
                                'caption': parts[1]
                            })
                self.data = pd.DataFrame(self.data)
        else:
            # 创建示例数据
            self.data = pd.DataFrame({
                'image_name': ['sample1.jpg', 'sample2.jpg'],
                'caption': ['a photo of a cat', 'a photo of a dog']
            })
    
    def __len__(self):
        return len(self.data)
    
    def __getitem__(self, idx):
        """获取单个样本"""
        # 获取图像文件名和描述
        image_name = self.data.iloc[idx]['image_name']
        caption = self.data.iloc[idx]['caption']
        
        # 加载图像
        image_path = os.path.join(self.image_dir, image_name)
        
        # 如果图像文件不存在，创建示例图像
        if not os.path.exists(image_path):
            # 创建示例图像（灰色图像）
            image = Image.new('RGB', (224, 224), color='gray')
        else:
            image = Image.open(image_path).convert('RGB')
        
        # 应用图像变换
        if self.transform:
            image = self.transform(image)
        
        # 使用processor处理图像和文本
        inputs = self.processor(
            text=[caption],
            images=image,
            return_tensors="pt",
            padding='max_length',  # 确保所有文本填充到相同长度
            truncation=True,
            max_length=self.max_length
        )
        
        # 移除batch维度
        pixel_values = inputs['pixel_values'].squeeze(0)
        input_ids = inputs['input_ids'].squeeze(0)
        attention_mask = inputs['attention_mask'].squeeze(0)
        
        return {
            'pixel_values': pixel_values,
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'caption': caption,
            'image_name': image_name
        }

class CIFAR10Dataset(Dataset):
    """CIFAR-10数据集类，专门用于CLIP模型训练"""
    
    def __init__(self, root, processor, train=True, download=True, transform=None, max_length=77):
        """
        初始化CIFAR-10数据集
        
        Args:
            root: 数据集保存路径
            processor: CLIP处理器
            train: 是否为训练集
            download: 是否下载数据集
            transform: 图像变换
            max_length: 文本最大长度
        """
        self.root = root
        self.processor = processor
        self.train = train
        self.transform = transform
        self.max_length = max_length
        
        # CIFAR-10类别名称
        self.class_names = ['airplane', 'automobile', 'bird', 'cat', 'deer', 
                           'dog', 'frog', 'horse', 'ship', 'truck']
        
        # 加载CIFAR-10数据集
        self.cifar_dataset = datasets.CIFAR10(
            root=root,
            train=train,
            download=download,
            transform=transforms.ToTensor()  # 基础变换，后面会处理
        )
        
        # 如果没有提供变换，使用默认变换
        if self.transform is None:
            self.transform = transforms.Compose([
                transforms.Resize((224, 224)),  # 调整大小以匹配CLIP输入
                transforms.ToTensor(),
                transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))
            ])
    
    def __len__(self):
        return len(self.cifar_dataset)
    
    def __getitem__(self, idx):
        """获取单个样本"""
        # 从CIFAR-10数据集获取图像和标签
        image_tensor, label_idx = self.cifar_dataset[idx]
        
        # 获取类别名称
        class_name = self.class_names[label_idx]
        
        # 创建文本描述
        caption = f"a photo of a {class_name}"
        
        # 将tensor转换为PIL图像
        image = transforms.ToPILImage()(image_tensor)
        
        # 应用变换
        if self.transform:
            image = self.transform(image)
        
        # 将图像值范围从[-1, 1]转换为[0, 1]，因为processor需要PIL图像格式
        if isinstance(image, torch.Tensor):
            # 对于已经归一化的tensor图像，先转换到[0, 1]范围
            image = (image + 1.0) / 2.0  # 从[-1, 1]转换到[0, 1]
            image = transforms.ToPILImage()(image)
        
        # 使用processor处理图像和文本
        inputs = self.processor(
            text=[caption],
            images=image,
            return_tensors="pt",
            padding='max_length',
            truncation=True,
            max_length=self.max_length
        )
        
        # 移除batch维度
        pixel_values = inputs['pixel_values'].squeeze(0)
        input_ids = inputs['input_ids'].squeeze(0)
        attention_mask = inputs['attention_mask'].squeeze(0)
        
        return {
            'pixel_values': pixel_values,
            'input_ids': input_ids,
            'attention_mask': attention_mask,
            'caption': caption,
            'label': label_idx,
            'class_name': class_name
        }

class ContrastiveDataset(Dataset):
    """对比学习数据集（包含正负样本对）"""
    
    def __init__(self, image_dir, text_file, processor, transform=None, num_negatives=4):
        """
        初始化对比学习数据集
        
        Args:
            num_negatives: 每个正样本的负样本数量
        """
        self.base_dataset = CLIPDataset(image_dir, text_file, processor, transform)
        self.num_negatives = num_negatives
        
    def __len__(self):
        return len(self.base_dataset)
    
    def __getitem__(self, idx):
        """获取样本，包含正样本和负样本"""
        # 正样本
        positive_sample = self.base_dataset[idx]
        
        # 负样本（随机选择）
        negative_indices = torch.randint(0, len(self.base_dataset), (self.num_negatives,))
        negative_samples = [self.base_dataset[i] for i in negative_indices]
        
        return {
            'positive': positive_sample,
            'negatives': negative_samples
        }

def create_sample_data(image_dir, text_file):
    """创建更好的示例数据"""
    # 创建图像目录
    os.makedirs(image_dir, exist_ok=True)
    
    # 更好的示例数据 - 使用更有意义的图像和描述
    sample_images = [
        ('animal_cat.jpg', 'a cute cat sitting on a sofa'),
        ('animal_dog.jpg', 'a happy dog playing in the park'),
        ('vehicle_car.jpg', 'a red car driving on the road'),
        ('building_house.jpg', 'a modern house with garden'),
        ('nature_tree.jpg', 'a tall tree in the forest'),
        ('food_pizza.jpg', 'a delicious pizza with toppings'),
        ('person_reading.jpg', 'a person reading a book'),
        ('city_street.jpg', 'a busy city street with buildings'),
        ('beach_sunset.jpg', 'a beautiful sunset at the beach'),
        ('mountain_view.jpg', 'a scenic mountain landscape')
    ]
    
    # 创建更有意义的示例图像
    for image_name, caption in sample_images:
        image_path = os.path.join(image_dir, image_name)
        if not os.path.exists(image_path):
            # 创建更有意义的图像（基于描述）
            if 'cat' in caption:
                # 创建猫的图像（橙色背景，黑色形状）
                image = Image.new('RGB', (224, 224), color=(255, 165, 0))
                # 添加简单的猫形状
                draw = ImageDraw.Draw(image)
                draw.ellipse([50, 50, 150, 150], fill=(0, 0, 0))  # 猫头
                draw.ellipse([60, 60, 80, 80], fill=(255, 255, 255))  # 眼睛
                draw.ellipse([120, 60, 140, 80], fill=(255, 255, 255))  # 眼睛
            elif 'dog' in caption:
                # 创建狗的图像（绿色背景，棕色形状）
                image = Image.new('RGB', (224, 224), color=(144, 238, 144))
                draw = ImageDraw.Draw(image)
                draw.rectangle([70, 70, 154, 154], fill=(139, 69, 19))  # 狗身体
                draw.ellipse([80, 50, 120, 90], fill=(139, 69, 19))  # 狗头
            elif 'car' in caption:
                # 创建汽车的图像（蓝色背景，红色汽车）
                image = Image.new('RGB', (224, 224), color=(173, 216, 230))
                draw = ImageDraw.Draw(image)
                draw.rectangle([60, 120, 164, 160], fill=(255, 0, 0))  # 汽车主体
                draw.rectangle([80, 100, 144, 120], fill=(255, 0, 0))  # 汽车顶部
            elif 'house' in caption:
                # 创建房子的图像（天空背景）
                image = Image.new('RGB', (224, 224), color=(135, 206, 235))
                draw = ImageDraw.Draw(image)
                draw.rectangle([80, 120, 144, 180], fill=(210, 180, 140))  # 房子主体
                draw.polygon([80, 120, 112, 80, 144, 120], fill=(139, 0, 0))  # 屋顶
            elif 'tree' in caption:
                # 创建树的图像（绿色背景）
                image = Image.new('RGB', (224, 224), color=(34, 139, 34))
                draw = ImageDraw.Draw(image)
                draw.rectangle([105, 140, 119, 180], fill=(101, 67, 33))  # 树干
                draw.ellipse([80, 100, 144, 140], fill=(0, 100, 0))  # 树冠
            elif 'pizza' in caption:
                # 创建披萨的图像（黄色背景）
                image = Image.new('RGB', (224, 224), color=(255, 255, 0))
                draw = ImageDraw.Draw(image)
                draw.ellipse([60, 60, 164, 164], fill=(255, 165, 0))  # 披萨
                # 添加配料
                for x, y in [(90, 90), (120, 100), (100, 130), (130, 80)]:
                    draw.ellipse([x, y, x+10, y+10], fill=(255, 0, 0))  # 番茄
            elif 'person' in caption:
                # 创建人物的图像
                image = Image.new('RGB', (224, 224), color=(240, 248, 255))
                draw = ImageDraw.Draw(image)
                draw.ellipse([100, 80, 124, 104], fill=(255, 218, 185))  # 头
                draw.rectangle([108, 104, 116, 140], fill=(0, 0, 255))  # 身体
                draw.rectangle([100, 140, 108, 160], fill=(0, 0, 255))  # 腿
                draw.rectangle([116, 140, 124, 160], fill=(0, 0, 255))  # 腿
            elif 'city' in caption:
                # 创建城市街景
                image = Image.new('RGB', (224, 224), color=(192, 192, 192))
                draw = ImageDraw.Draw(image)
                # 建筑物
                for x in [40, 80, 120, 160]:
                    height = torch.randint(80, 150, (1,)).item()
                    draw.rectangle([x, 224-height, x+30, 224], fill=(105, 105, 105))
            elif 'beach' in caption:
                # 创建海滩日落
                image = Image.new('RGB', (224, 224), color=(255, 165, 0))
                draw = ImageDraw.Draw(image)
                # 海洋
                draw.rectangle([0, 150, 224, 224], fill=(0, 0, 139))
                # 太阳
                draw.ellipse([100, 80, 124, 104], fill=(255, 255, 0))
            elif 'mountain' in caption:
                # 创建山脉景观
                image = Image.new('RGB', (224, 224), color=(135, 206, 235))
                draw = ImageDraw.Draw(image)
                # 山脉
                draw.polygon([0, 224, 80, 120, 160, 180, 224, 224], fill=(139, 69, 19))
                draw.polygon([80, 224, 140, 100, 224, 224], fill=(101, 67, 33))
            else:
                # 默认创建彩色图像
                image = Image.new('RGB', (224, 224), color=(
                    torch.randint(50, 200, (1,)).item(),
                    torch.randint(50, 200, (1,)).item(),
                    torch.randint(50, 200, (1,)).item()
                ))
            
            image.save(image_path)
    
    # 创建文本描述文件
    with open(text_file, 'w', encoding='utf-8') as f:
        for image_name, caption in sample_images:
            f.write(f"{image_name}\t{caption}\n")
    
    print(f"更好的示例数据已创建: {len(sample_images)} 个样本")
    print("图像描述包含更丰富的语义信息，能够更好地测试CLIP模型")
    print("注意：这些是简单的几何图像，对于真实的CLIP模型测试，建议使用真实图像数据")
    
    # 提供下载真实示例数据的选项
    print("\n要获得更好的测试效果，可以下载真实图像数据:")
    print("1. 使用在线图像数据集（如COCO、Flickr等）")
    print("2. 使用torchvision.datasets中的预定义数据集")
    print("3. 使用网络爬虫获取相关图像")
    
    return sample_images